// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: xguo

#pragma once

#include <algorithm>
#include <cstdint>
#include <deque>
#include <memory>
#include <string>
#include <utility>

#include "coin2/base/proto_util.h"
#include "coin2/exchange/base/api_base/asio_http_context.h"
#include "coin2/exchange/base/order/order_context_manager.h"

namespace coin::proto {
class RateLimitReportMessage;
}

namespace coin2::exchange::base::order {

class IRateTracker {
 public:
  using ActionType = coin::proto::RequestType;
  using RateLimitReportMessage = coin::proto::RateLimitReportMessage;
  using OrderLoggerCB = std::function<void(const RateLimitReportMessage&, bool, int64_t)>;
  using HttpContextPtr = std::shared_ptr<AsioHttpContext>;

 public:
  virtual ~IRateTracker() {}
  virtual void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) = 0;
  virtual void Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) = 0;
  virtual bool HasEnoughRateLimit(int64_t timestamp, ActionType action_type) = 0;
  virtual void ParseRateLimitInfo(const HttpContextPtr& http_context, ActionType action_type) = 0;
  virtual void SetRateLimitInfoString(const google::protobuf::Message* message) = 0;
  virtual void WriteRateLimitReportMessage(OrderLoggerCB callback) = 0;
  virtual double GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) = 0;
};

class RateTracker {
 public:
  struct Entry {
    Entry(int32_t num_reqs, int64_t timestamp) : num_reqs(num_reqs), timestamp(timestamp) {}
    int32_t num_reqs;
    int64_t timestamp;
  };

 public:
  RateTracker() = default;
  explicit RateTracker(int64_t window_size) : window_size_(window_size) {}
  void set_window_size(int64_t window_size) { window_size_ = window_size; }

  int64_t MaxReq() const { return rate_limit_; }
  int64_t GetRemaining(int64_t timestamp) {
    UpdateTime(timestamp);
    int64_t last_update_time = timestamp - window_size_;
    if (!queue_.empty()) {
      last_update_time = queue_.back().timestamp;
    }
    const double delta = (timestamp - last_update_time) / 1.0e9;
    const double wsize = window_size_ / 1.0e9;
    const int64_t added = delta / wsize * rate_limit_;
    DCHECK_GE(added, 0) << name_ << ": " << timestamp << ", " << last_update_time << ", "
                        << rate_limit_ << ", " << window_size_;

    const int64_t estimated_remaining = std::min(rate_limit_, rate_limit_remaining_ + added);
    DLOG(INFO) << name_ << ": "
               << "rate_limit: " << rate_limit_ << ", estimated remaining: " << estimated_remaining
               << ", safety_left: " << safety_left_ << ", sum: " << sum_
               << ", rate_limit_remaining: " << rate_limit_remaining_;
    return std::max<int64_t>(0, estimated_remaining);
  }

  void UpdateTime(int64_t timestamp) {
    while (!queue_.empty() && (timestamp - queue_.at(0).timestamp > window_size_)) {
      sum_ -= queue_.at(0).num_reqs;
      rate_limit_remaining_ += queue_.at(0).num_reqs;
      queue_.pop_front();
    }
  }

  int64_t GetUsedRate(int64_t timestamp) {
    UpdateTime(timestamp);
    return sum_;
  }

  bool HasEnoughRateLimit(int64_t timestamp) {
    UpdateTime(timestamp);
    int64_t last_update_time = timestamp - window_size_;
    if (!queue_.empty()) {
      last_update_time = queue_.back().timestamp;
    }
    const double delta = (timestamp - last_update_time) / 1.0e9;
    const double wsize = window_size_ / 1.0e9;
    const int64_t added = delta / wsize * rate_limit_;
    DCHECK_GE(added, 0) << name_ << ": " << timestamp << ", " << last_update_time << ", "
                        << rate_limit_ << ", " << window_size_;
    const int64_t estimated_remaining = std::min(rate_limit_, rate_limit_remaining_ + added);
    DLOG(INFO) << name_ << ": "
               << "rate_limit: " << rate_limit_ << ", estimated remaining: " << estimated_remaining
               << ", safety_left: " << safety_left_ << ", sum: " << sum_
               << ", rate_limit_remaining: " << rate_limit_remaining_;
    return estimated_remaining > safety_left_ && sum_ < rate_limit_ - safety_left_;
  }

  bool HasEnoughRateLimitNoSafe(int64_t timestamp) {
    UpdateTime(timestamp);
    int64_t last_update_time = timestamp - window_size_;
    if (!queue_.empty()) {
      last_update_time = queue_.back().timestamp;
    }
    const double delta = (timestamp - last_update_time) / 1.0e9;
    const double wsize = window_size_ / 1.0e9;
    const int64_t added = delta / wsize * rate_limit_;
    DCHECK_GE(added, 0) << name_ << ": " << timestamp << ", " << last_update_time << ", "
                        << rate_limit_ << ", " << window_size_;
    const int64_t estimated_remaining = std::min(rate_limit_, rate_limit_remaining_ + added);
    DLOG(INFO) << name_ << ": "
               << "rate_limit: " << rate_limit_ << ", estimated remaining: " << estimated_remaining
               << ", safety_left: " << safety_left_ << ", sum: " << sum_
               << ", rate_limit_remaining: " << rate_limit_remaining_;
    return estimated_remaining > 0;
  }

  // not using
  bool HasEnoughOrderRateLimit(int64_t timestamp) {
    UpdateTime(timestamp);
    int64_t last_update_time = timestamp - window_size_;
    if (!queue_.empty()) {
      last_update_time = queue_.back().timestamp;
    }
    const double delta = (timestamp - last_update_time) / 1.0e9;
    const double wsize = window_size_ / 1.0e9;
    const int64_t added = delta / wsize * rate_limit_;
    const int64_t estimated_remaining = std::min(rate_limit_, rate_limit_remaining_ + added);
    DLOG(INFO) << name_ << ": "
               << "estimated remaining: " << estimated_remaining << ", sum: " << sum_
               << ", rate_limit_remaining: " << rate_limit_remaining_;
    return estimated_remaining > 5;
  }

  void Push(int32_t num_reqs, int64_t timestamp) {
    queue_.emplace_back(num_reqs, timestamp);
    sum_ += num_reqs;
    rate_limit_remaining_ -= num_reqs;
    UpdateTime(timestamp);
  }

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) {
    /*
     "x-ratelimit-limit": 60
     "x-ratelimit-remaining": 58
     "x-ratelimit-reset": 1489791662
     */
    CHECK(rate_limit > 0) << "invalid rate limit: " << rate_limit;
    CHECK(rate_limit_remaining >= 0) << "invalid rate limit remaining: " << rate_limit_remaining;
    rate_limit_ = rate_limit;
    rate_limit_remaining_ = rate_limit_remaining;
    sum_ = rate_limit - rate_limit_remaining;
    safety_left_ = std::min(rate_limit_ / 6, 30L);
  }

  void SetSafetyLeft(int32_t safety_left) { safety_left_ = safety_left; }

  void SetName(const std::string name) { name_ = name; }
  auto GetName() const -> const std::string& { return name_; }

 private:
  int64_t window_size_;
  int64_t sum_{0};
  std::deque<Entry> queue_;
  int64_t rate_limit_{60};
  int64_t rate_limit_remaining_{60};
  int64_t safety_left_{10};
  std::string name_{"rate_tracker"};
};

class RateLimitReportWriter {
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using OrderLoggerCB = coin2::exchange::base::order::IRateTracker::OrderLoggerCB;
  using RateLimitReportMessage = coin2::exchange::base::order::IRateTracker::RateLimitReportMessage;

  class RateLimitReportWriterImpl {
   public:
    explicit RateLimitReportWriterImpl(ActionType action_type)
        : action_type_{action_type}, num_total_hit_{0}, last_write_ts_{0} {}

    void
    Write(OrderLoggerCB callback, RateLimitReportMessage* rate_limit_report_msg, int64_t interval) {
      if (rate_limit_report_msg->last_hit_request_type() != action_type_) {
        return;
      }
      ++num_total_hit_;
      auto now = GetCurrentTimestamp();
      if (now - last_write_ts_ < interval) {
        return;
      }

      last_write_ts_ = now;
      rate_limit_report_msg->set_num_total_hit(num_total_hit_);
      if (callback) {
        callback(*rate_limit_report_msg, true, now);
      }
    }

   private:
    ActionType action_type_;
    int64_t num_total_hit_;
    int64_t last_write_ts_;
  };

 public:
  RateLimitReportWriter()
      : rate_limit_report_msg_{std::make_unique<RateLimitReportMessage>()},
        query_rlr_writer_{ActionType::QUERY},
        place_rlr_writer_{ActionType::PLACE},
        cancel_rlr_writer_{ActionType::CANCEL},
        amend_rlr_writer_{ActionType::AMEND} {
    rate_limit_report_msg_->set_rate_limit_info("");
    rate_limit_report_msg_->set_num_total_hit(0);
    rate_limit_report_msg_->set_last_hit_request_type(ActionType::UNKNOWN);
    rate_limit_report_msg_->set_last_hit_timestamp(0);
  }

  void Write(OrderLoggerCB callback) {
    if (!callback) {
      return;
    }
    auto report_msg = rate_limit_report_msg_.get();
    switch (report_msg->last_hit_request_type()) {
      case ActionType::PLACE:
        place_rlr_writer_.Write(callback, report_msg, interval_);
        break;
      case ActionType::AMEND:
        amend_rlr_writer_.Write(callback, report_msg, interval_);
        break;
      case ActionType::CANCEL:
        cancel_rlr_writer_.Write(callback, report_msg, interval_);
        break;
      case ActionType::QUERY:
        query_rlr_writer_.Write(callback, report_msg, interval_);
        break;
      default:
        // shouldn't reach here
        break;
    }
  }

  void SetHitRateLimit(int64_t timestamp, ActionType action_type) {
    rate_limit_report_msg_->set_last_hit_request_type(action_type);
    rate_limit_report_msg_->set_last_hit_timestamp(timestamp);
  }

  void SetRateLimitInfoString(const google::protobuf::Message* message) {
    rate_limit_report_msg_->set_rate_limit_info(ProtoToJsonString(*message));
  }

 private:
  std::unique_ptr<RateLimitReportMessage> rate_limit_report_msg_ = {};
  RateLimitReportWriterImpl query_rlr_writer_;
  RateLimitReportWriterImpl place_rlr_writer_;
  RateLimitReportWriterImpl cancel_rlr_writer_;
  RateLimitReportWriterImpl amend_rlr_writer_;
  const int64_t interval_ = 5'000'000'000LL;  // 5 seconds
};

class DefaultRateTracker : public coin2::exchange::base::order::IRateTracker {
 public:
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;

 public:
  explicit DefaultRateTracker(int64_t window_size)
      : rate_tracker_{std::make_unique<RateTracker>(window_size)},
        rate_limit_report_writer_{std::make_unique<RateLimitReportWriter>()} {}

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) override {
    CHECK(rate_limit > 0) << "invalid rate limit setting: " << rate_limit;
    rate_tracker_->SetRateLimitInfo(rate_limit, rate_limit_remaining);
  }

  void Push(int32_t num_reqs, int64_t timestamp, [[maybe_unused]] ActionType action_type) override {
    rate_tracker_->Push(num_reqs, timestamp);
  }

  bool HasEnoughRateLimit(int64_t timestamp, [[maybe_unused]] ActionType action_type) override {
    auto ret = rate_tracker_->HasEnoughRateLimit(timestamp);
    if (!ret) {
      rate_limit_report_writer_->SetHitRateLimit(timestamp, action_type);
    }
    return ret;
  }

  void ParseRateLimitInfo(
      [[maybe_unused]] const std::shared_ptr<AsioHttpContext>& http_context,
      [[maybe_unused]] ActionType action_type) override {}

  void WriteRateLimitReportMessage(OrderLoggerCB callback) override {
    if (!callback) {
      return;
    }
    rate_limit_report_writer_->Write(callback);
  }

  void SetRateLimitInfoString(const google::protobuf::Message* message) override {
    rate_limit_report_writer_->SetRateLimitInfoString(message);
  }

  double GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) override {
    if (!rate_tracker_) {
      return 0.;
    }
    auto max_reqs = rate_tracker_->MaxReq();
    auto remaining = rate_tracker_->GetRemaining(timestamp);
    if (max_reqs < remaining) {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, rate_limit=" << max_reqs
                 << ", remaining=" << remaining;
      return 0.;
    }
    return 1.0 * (max_reqs - remaining) / max_reqs;
  };

 private:
  std::unique_ptr<RateTracker> rate_tracker_;
  std::unique_ptr<RateLimitReportWriter> rate_limit_report_writer_;
};

class FixedWindowRateTracker {
 public:
  FixedWindowRateTracker() = default;
  explicit FixedWindowRateTracker(int64_t window_size) : window_size_(window_size) {}
  void set_window_size(int64_t window_size) { window_size_ = window_size; }
  void set_offset(int64_t offset) { offset_ = offset; }

  int64_t MaxReq() const { return rate_limit_; }
  int64_t GetRemaining(int64_t timestamp) {
    UpdateTime(timestamp);
    return std::max<int64_t>(0, rate_limit_ - sum_);
  }

  void UpdateTime(int64_t timestamp) {
    int64_t window_index = (timestamp - offset_) / window_size_;
    if (window_index > current_window_index_) {
      current_window_index_ = window_index;
      sum_ = 0;
    }
  }

  int64_t GetUsedRate(int64_t timestamp) {
    UpdateTime(timestamp);
    return sum_;
  }

  void Push(int32_t num_reqs, int64_t timestamp) {
    UpdateTime(timestamp);
    sum_ += num_reqs;
  }

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) {
    rate_limit_ = rate_limit;
    sum_ = rate_limit - rate_limit_remaining;
  }

  void SetName(const std::string name) { name_ = name; }
  auto GetName() const -> const std::string& { return name_; }

 private:
  int64_t window_size_{0};
  int64_t current_window_index_{0};
  int64_t sum_{0};
  int64_t rate_limit_{60};
  int64_t offset_{0};
  std::string name_{"rate_tracker"};
};

class OrderRateTracker {
 public:
  using OrderEvent = coin::proto::OrderEvent;

  explicit OrderRateTracker(const MarketExchangeApi& mea) : mea_(mea) {
    // config read from og_info.json
    SetOrderRatioLimitWindow(mea);
  }

  void SetOrderRatioLimitWindow(const MarketExchangeApi& mea) {
    if (mea.exchange == ExchangeType::Huobi && mea.market == MarketType::Futures) {
      submit_order_rate_tracker_ = std::make_unique<RateTracker>(603'000'000'000LL);
      accept_order_rate_tracker_ = std::make_unique<RateTracker>(603'000'000'000LL);
      cancel_order_rate_tracker_ = std::make_unique<RateTracker>(603'000'000'000LL);
    } else {
      submit_order_rate_tracker_ = std::make_unique<RateTracker>(1'000'000'000LL);
      accept_order_rate_tracker_ = std::make_unique<RateTracker>(1'000'000'000LL);
      cancel_order_rate_tracker_ = std::make_unique<RateTracker>(1'000'000'000LL);
    }
  }

  int64_t GetRecentSubmittedOrderCount(int64_t ts) const {
    return submit_order_rate_tracker_->GetUsedRate(ts);
  }
  int64_t GetRecentAcceptedOrderCount(int64_t ts) const {
    return accept_order_rate_tracker_->GetUsedRate(ts);
  }
  int64_t GetRecentCancelledOrderCount(int64_t ts) const {
    return cancel_order_rate_tracker_->GetUsedRate(ts);
  }

  void HandleOrderEvent(const OrderEvent& order_event, const OrderContext* oc, int64_t ts) {
    if (order_event.type() == OrderEvent::ORDER_SUBMITTED) {
      submit_order_rate_tracker_->Push(1, ts);
    } else if (order_event.type() == OrderEvent::ORDER_ACCEPTED) {
      accept_order_rate_tracker_->Push(1, ts);
    } else if (
        order_event.type() == OrderEvent::ORDER_AUTO_CANCELED ||
        order_event.type() == OrderEvent::CANCEL_CONFIRMED) {
      if (oc->filled_qty() < 1e-10 &&
          oc->cancel_confirmed_time() - oc->order_accepted_time() < 3e9) {
        cancel_order_rate_tracker_->Push(1, ts);
      }
    }
  }

  bool HasEnoughRateLimit(int64_t ts) const {
    if (mea_.exchange == ExchangeType::Huobi && mea_.market == MarketType::Futures) {
      int total_submit_num = submit_order_rate_tracker_->GetUsedRate(ts);
      int total_accept_num = accept_order_rate_tracker_->GetUsedRate(ts);
      int total_cancel_num = cancel_order_rate_tracker_->GetUsedRate(ts);
      if (total_accept_num >= 2400 && total_cancel_num >= total_accept_num * 0.98) {
        LOG(INFO) << "[CancelRateLimit] Total sumbitted: " << total_submit_num
                  << " Total accepted: " << total_accept_num
                  << " Total canceled: " << total_cancel_num;
        return false;
      }
    }
    return true;
  }

 private:
  const MarketExchangeApi mea_;
  std::unique_ptr<RateTracker> submit_order_rate_tracker_;
  std::unique_ptr<RateTracker> accept_order_rate_tracker_;
  std::unique_ptr<RateTracker> cancel_order_rate_tracker_;
};

}  // namespace coin2::exchange::base::order
